#! /usr/bin/python

""" This script does a spectral analysis of a single data set and returns the power spectral density before, during and after a stimulus is presented. In addition a spectrogram is provided """

import os
import sys
from collections import defaultdict

import numpy
import scipy
from matplotlib import pyplot

from LFP import signal_utils
from projects.electrophysiology.view_data import file_utils

# SPECGRAM VARIABLES
time_resolution = 0.005 # in seconds
frequency_resolution = 1 # in Hz
set_xlim = True # set to False to plot entire time range.
xlim_low = 30 #in sec
xlim_high = 40 #in sec
low_frequency = 10
high_frequency = 100

bath_filenames = []
signal_filenames = defaultdict(list)

for arg in sys.argv[1:]:
    if 'bath=' in arg:
        bath_filenames.append(arg.split('=')[1])
    elif '=' in arg:
        figure_name, filename = arg.split('=')
        signal_filenames[figure_name].append(filename)
    else: 
        signal_filenames['default'].append(arg)

# First extract the data

def extract_data(filename):
    return file_utils.read_wessel_text_file(filename) 

def get_voltages(filenames):
    voltages = []
    sampling_freq = None
    pulse_index = None
    for filename in filenames:
        print "Reading in %s" % filename
        voltage, sampling_freq, pulse_index = extract_data(filename)
        voltages.append(voltage)
    return voltages, sampling_freq, pulse_index
    
signals = {}
for figure_name, filenames in signal_filenames.items():
    signal_voltages, sampling_freq, pulse_index = get_voltages(filenames)
    signals[figure_name] = (signal_voltages, sampling_freq, pulse_index)
 
bath_voltages = get_voltages(bath_filenames)[0]

max_Pxx = 0.0
min_Pxx = None
Pxx_results = {}
for figure_name, (signal_voltages, 
                    sampling_freq, 
                       pulse_index) in signals.items():
    print "Computing the average specgram."
    Pxx, freqs, bins = signal_utils.avg_specgram(signal_voltages,
                                 sampling_freq, time_resolution, 
                                 frequency_resolution, 
                                 high_frequency_cutoff = 2*high_frequency+100,
                                 bath_signals=bath_voltages)
    print "Found minimum power for %s to be %f" % (figure_name, numpy.min(Pxx))
    print "Found max power for %s to be %f" % (figure_name, numpy.max(Pxx))
    max_Pxx = max(max_Pxx, numpy.max(Pxx))
    if min_Pxx is None:
        min_Pxx = numpy.min(Pxx)
    min_Pxx = min(min_Pxx, numpy.min(Pxx))
    Pxx_results[figure_name] = (Pxx, freqs, bins)

print "Found overall Pxx min to be: %f" % min_Pxx
print "Found overall Pxx max to be: %f" % max_Pxx
    
for figure_name, [Pxx, freqs, bins] in Pxx_results.items():
    fig = pyplot.figure()
    top_axes = fig.add_subplot(211)
    pyplot.title('%s' % figure_name)
    bottom_axes = fig.add_subplot(212, sharex=top_axes, sharey=top_axes)
    Pxx[-1,-1] = max_Pxx
    Pxx[-1,-2] = min_Pxx
    print ("Finally minimum power for %s to be %f (%f)" % 
             (figure_name, numpy.min(Pxx), min_Pxx))
    print ("Finally max power for %s to be %f (%f)" % 
             (figure_name, numpy.max(Pxx), max_Pxx))
    signal_utils.plot_specgram(Pxx, freqs, bins, top_axes,
                                            logscale=True)
    top_axes.set_ylim(low_frequency,high_frequency)
    signal_utils.plot_specgram(Pxx, freqs, bins, bottom_axes,
                                            logscale=False)
    bottom_axes.set_ylim(low_frequency,high_frequency)
    if set_xlim:
        bottom_axes.set_xlim(xlim_low,xlim_high)

pyplot.show()





